import numpy as np
import matplotlib.pyplot as plt
import math
from copy import deepcopy

np.set_printoptions(precision=3, suppress=True, linewidth=100)


class Matrix_Solutions(object):
    """Matrix_Solutions by Junno
    Args:
        object ([type]): [description]
    """

    def __init__(self,):
        return

    @classmethod
    def Full_Rank_Fact(self, target, test=False):
        """Full Rank Decomposition Solution for 2D-matrix by Junno

        Args:
            target ([np.darray]): [matrix to be processed]
            test (bool, optional): [whether to show testing information]. Defaults to False.
        """
        assert len(np.shape(target)) == 2
        # get row ans col numbers
        M, N = np.shape(target)

        # A for processing and raw_A for result extraction
        raw_A = deepcopy(target)
        A = deepcopy(target)
        if test:
            print('original matrix:')
            print(target)

        for i in range(M):
            for j in range(i, N):
                if test:
                    print('During process in row:{}, col:{}'.format(i, j))
                if sum(abs(A[i:, j])) != 0:
                    zero_index = []
                    non_zero_index = []
                    for k in range(i, M):
                        if abs(A[k, j]):
                            non_zero_index.append(k)
                        else:
                            zero_index.append(k)

                    non_zero_index = sorted(
                        non_zero_index, key=lambda x: abs(A[x, j]))
                    sort_cols = non_zero_index+zero_index
                    if test:
                        print("Sort_cols index: {}".format(sort_cols))

                    A[i:, :] = A[sort_cols, :]

                    if test:
                        print('After resorting')
                        print(A)

                    prefix = -A[i+1:, j]/A[i, j]
                    temp = (np.array(prefix).reshape(-1, 1)
                            )@A[i, :].reshape((1, -1))
                    A[i+1:, :] += temp

                    for k in range(i):
                        if A[k, j] != 0:
                            A[k, :] += -(A[k, j]/A[i, j])*A[i, :]

                    if test:
                        print('After primary row transformation:')
                        print(A)

                    break

                else:
                    continue

        # list to record indexs of principal maximum
        principal_indexs = [[], []]
        for m in range(M):
            for n in range(m, N):
                if A[m, n] != 0:
                    principal_indexs[0].append(n)
                    principal_indexs[1].append(m)
                    if A[m, n] != 1:
                        A[m, :] *= 1./A[m, n]
                    break
        # -0 to 0
        A[abs(A) == 0] = 0

        if test:
            print('final result')
            print(A)
            print('principal maximum index: {}'.format(principal_indexs))

        # choose principal maximum and return B and C
        if len(principal_indexs[0]):
            B = raw_A[:, principal_indexs[0]]
            C = A[principal_indexs[1], :]
            return B, C
        else:
            raise ValueError("Can't find any principle maximum @_@")

    @classmethod
    def LDV_2D(self, A, eps=1e-6, mode='LDV', Test=False, solve_equation=False):
        """custom function for 2D-matrix LDV-Decomposition by Junno

        Args:
            A ([np.darray]): [2D-matrix to be processed with LDV-Decomposition ]
            mode (str, optional): [choose mode]. Defaults to LDV.
            eps ([float]): numerical threshold.
            Test (bool, optional): [show checking information]. Defaults to False.
            solve_equation(bool,optional): [True when use this function to solve equation like Ax=b]. Defaults to True.

        Raises:
            ValueError: [matrix don't meet the requirement for LDV or LU decomposition]

        Returns:
            [np.darray]: [resort-matrix, L, D, V] or [resort-matrix, L, U]
        """
        assert len(A.shape) == 2 and A.shape[0] == A.shape[1]
        # primary check for LDV:
        if mode == 'LDV':
            if not self.Check_full_rank(A):
                raise ValueError(
                    "This matrix can't be operated with LDV decomposition, maybe you can try LU decomposition instead ^_^")
        M, N = A.shape
        cm = np.copy(A)
        resort_m = np.copy(A)
        operator_list = []
        L = np.eye(M)
        #print('original matrix:\n', A)

        for i in range(M):
            for j in range(i, N):
                if Test:
                    print('During process in row:{}, col:{}'.format(i, j))
                if sum(abs(cm[i:, j])) > eps:
                    zero_index = []
                    non_zero_index = []
                    for k in range(i, M):
                        if abs(cm[k, j]) > eps:
                            non_zero_index.append(k)
                        else:
                            zero_index.append(k)

                    non_zero_index = sorted(
                        non_zero_index, key=lambda x: abs(cm[x, j]))
                    sort_idnex = non_zero_index+zero_index

                    if Test:
                        print('Sorting index for {} cols:\n'.format(i), sort_idnex)

                    if sort_idnex[0] != i:
                        operator_list.append(['t', [sort_idnex[0], i]])
                        cm[[sort_idnex[0], i], :] = cm[[i, sort_idnex[0]], :]
                        resort_m[[sort_idnex[0], i],
                                 :] = resort_m[[i, sort_idnex[0]], :]
                    if Test:
                        print('After resorting\n', cm)

                    if np.sum(abs(cm[i+1:, j])) > eps:
                        prefix = -cm[i+1:, j]/cm[i, j]
                        operator_list.append(['pre', i, prefix.tolist()])

                        # operate elementary row transformation on matrix based on elimination coefficient
                        temp = (np.array(prefix).reshape(-1, 1)
                                )@cm[i, :].reshape((1, -1))
                        cm[i+1:, :] += temp

                    if Test:
                        print('After elementary row transformation:')
                        print(cm)

                    break

                else:
                    continue

        # Prints a new matrix arranged according to the main elements (just swap rows without changing the matrix information)
        if Test:
            print('resort matrix:\n', resort_m)
            print('After primary row transformation: \n', cm)
            print('Operation list:\n', operator_list)

        # Operate on the identity matrix according to the saved sorting and elimination coefficients to get the matrix L
        for i in range(len(operator_list)-1, -1, -1):
            op = operator_list[i]
            if op[0] == 't':
                m, n = op[1]
                L[[m, n], :] = L[[n, m], :]
            elif op[0] == 'pre':
                ind = op[1]
                temp = -(np.array(op[2]).reshape(-1, 1)
                         )@L[ind, :].reshape((1, -1))
                L[ind+1:, :] += temp
            elif op[0] == 'jump':
                ind = op[1]
                L[ind+1, :] -= op[2]*L[ind, :]

        # Operate the same row transpose on the L to let it be lower triangle
        for i in range(len(operator_list)):
            op = operator_list[i]
            if op[0] == 't':
                m, n = op[1]
                L[[m, n], :] = L[[n, m], :]

        # LU mode
        if mode == 'LU':
            if solve_equation:
                return operator_list, resort_m, L, cm
            else:
                return resort_m, L, cm
        # LDV mode
        elif mode == 'LDV':
            # extract the principal diagonal elements of the transformed matrix
            diag = np.diag(cm)

            # Error
            if 0 in diag:
                raise ValueError(
                    'this matrix can\'t be operated with LDV decomposition')
            return resort_m, L, np.diag(diag), np.triu(cm/diag.reshape((-1, 1)))

    @classmethod
    def Check_full_rank(self, target, eps=1e-6):
        """Check_full_rank Function

        Args:
            target ([np.darray]): input matrix

        Returns:
            [bool]: [whether full rank]

        Last edited: 22-01-02
        Author: Junno
        """
        if np.sum(target) < eps:
            # all-zero
            return False
        # get row ans col numbers
        M, N = np.shape(target)
        B, C = self.Full_Rank_Fact(target, False)
        return B.shape[1] == min(M, N)

    @classmethod
    def Flip(self, M):
        """Flip Matrix

        Args:
            M ([np.darray]): input matrix

        Returns:
            [np.darray]: [result]
        """
        assert len(M.shape) == 2
        if M.shape[1] > 1:
            return np.fliplr(np.flipud(M))
        else:
            return np.flipud(M)

    @staticmethod
    def Gaussian_eliminate(target, b, eps=1e-3, test=False):
        """Gaussian_eliminate to solve Ax=b equation

        Args:
            A ([np.darray]): [A]
            b ([np.darray]): [b]
            eps ([float]): numerical threshold
            test ([bool]): show checking information, default to False.

        Returns:
            [np.darray]: [the solution to equation]
        """
        M, N = target.shape
        A = np.concatenate((target, b), axis=1)
        for i in range(M):
            if test:
                print('During process in row:{}, col:{}'.format(i, i))
            if sum(abs(A[i:, i])) > eps:
                zero_index = []
                non_zero_index = []
                for k in range(i, M):
                    if abs(A[k, i]) > eps:
                        non_zero_index.append(k)
                    else:
                        zero_index.append(k)

                non_zero_index = sorted(
                    non_zero_index, key=lambda x: abs(A[x, i]))
                sort_cols = non_zero_index+zero_index
                if test:
                    print("Sort_cols index: {}".format(sort_cols))

                A[i:, :] = A[sort_cols, :]

                if test:
                    print('After resorting')
                    print(A)

                if np.sum(abs(A[i+1:, :])):
                    prefix = -A[i+1:, i]/A[i, i]
                    temp = (np.array(prefix).reshape(-1, 1)
                            )@A[i, :].reshape((1, -1))
                    A[i+1:, :] += temp
                    if test:
                        print('After PRM')
                        print(A)

        for m in range(M-1):
            for n in range(m, N):
                if abs(A[m, n]) > eps:
                    A[m, :] *= 1./A[m, n]
                    break

        if test:
            print(A)

        free_dims = []

        for i in range(M):
            if A[i, i] == 0:
                free_dims.append(i)

        if M in free_dims:
            print("A and b is unmatch, you can use LSM to solve this equation!")
        else:
            new_b = A[:, -1].reshape((M,))
            A = A[:, :-1]

            if test:
                print('free dims: ', free_dims)

            x = np.zeros((1, M), dtype=target.dtype)
            for i in range(M-1, -1, -1):
                if i in free_dims:
                    x[0, i] = 1
                    if np.sum(A[i, i:]*x[0, i:]) != new_b[i]:
                        print(
                            "A and b is unmatch, you can use the Least Square Method to solve this equation!")
                        return None
                else:
                    x[0, i] = (new_b[i]-np.sum(A[i, i+1:]*x[0, i+1:]))/A[i, i]

            if test:
                if len(free_dims):
                    print("Answer is not unique with free_dims: {}".format(*free_dims))
                else:
                    print("The given answer is unique")
            return x

    @classmethod
    def Solve_Ax_b_Equation(self, A, b_, mode='LU', test=False):
        """Function to Solve Ax_b_like Equation using LU decomposition by Junno

        Args:
            A ([np.darray]): [A]
            b ([np.darray]): [b]
            mode ([string]): LU for standard solver, LSM for Least-square-method for incompatible equations, Gaussian for the Gaussian-elimination-method
            test ([bool]): show checking information, default to False
        Returns:
            [np.darray]: answer [x]

        Last edited: 21-12-31
        Author: Junno
        """
        assert len(A.shape) == len(b.shape)
        assert A.shape[0] == b_.shape[0]
        b = np.copy(b_)
        if mode == "LU":
            # Check that whether the ranks of A and B are the same or not
            if self.Get_row_rank(A) != self.Get_row_rank(np.concatenate((A, b), axis=1)):
                raise ValueError(
                    "This matrix equation is incompatible, maybe you can try the Least Square Method ^_^")
            else:
                operator_list, trans_m, L, U = self.LDV_2D(
                    A, mode="LU", Test=False, solve_equation=True)
            #     print(operator_list)
                if test:
                    print("L matrix from LU:")
                    print(L)
                    print("U matrix from LU:")
                    print(U)
                # transform b according with L
                for i in range(len(operator_list)):
                    op = operator_list[i]
                    if op[0] == 't':
                        m, n = op[1]
                        b[[m, n], :] = b[[n, m], :]

                if test:
                    print('resort b: \n', b)

                Y = np.zeros_like(b)
                X = np.zeros_like(b)

                for i in range(L.shape[0]):
                    temp = 0.
                    for j in range(i):
                        temp += Y[j, 0]*L[i, j]
                    Y[i, 0] = (b[i, 0]-temp)/L[i, i]

                if test:
                    print("Calculate intermediate param Y")
                    print(Y)
                # filp U and Y
                flip_U = self.Flip(U)
                flip_Y = self.Flip(Y)

                if test:
                    print("Filp U:")
                    print(flip_U)
                    print("Filp Y:")
                    print(flip_Y)

                # check zero in diag(flip_U)
                if 0 in np.diag(flip_U):
                    if test:
                        print(np.diag(flip_U))
                    raise ValueError("Can't solve this equations")
                else:
                    for i in range(flip_U.shape[0]):
                        temp = 0.
                        for j in range(i):
                            temp += X[j, 0]*flip_U[i, j]
                        X[i, 0] = (flip_Y[i, 0]-temp)/flip_U[i, i]

                    return self.Flip(X)

        elif mode == "LSM":
            MP_A = self.Moore_Penrose_General_Inv(A)
            app_x = MP_A@b
            print('approximation error: ', np.linalg.norm(A@app_x-b))
            return app_x

        elif mode == "Gaussian":
            x = self.Gaussian_eliminate(A, b)
            return x

    @classmethod
    def Get_col_rank(self, target, eps=1e-6):
        """Get_col_rank of target

        Args:
            target ([np.darray]): input matrix

        Returns:
            [int]: [rank of cal]

        Last edited: 22-01-02
        Author: Junno
        """
        if np.sum(target) < eps:
            # all-zero
            return 0
        B, C = self.Full_Rank_Fact(target, False)
        return B.shape[1] == target.shape[1]

    @classmethod
    def Get_row_rank(self, target):
        """Get_row_rank of target

        Args:
            target ([np.darray]): input matrix

        Returns:
            [int]: [rank of row]

        Last edited: 22-01-02
        Author: Junno
        """
        if np.sum(target) < eps:
            # all-zero
            return 0
        B, C = self.Full_Rank_Fact(target, False)
        return B.shape[1] == target.shape[0]

    @classmethod
    def Primary_Row_Transformation_Inv(self, target, eps=1e-6, test=False):
        """Primary_Row_Transformation_Inv 

        Args:
            target ([np.darray]): [target matrix]
            eps ([float]): numerical threshold.
            test (bool, optional): [whether to show testing information]. Defaults to False.

        Returns:
            [np.darray]: [A^-]
        """
        assert len(np.shape(target)) == 2
        # get row and col numbers
        M, N = np.shape(target)
        assert M == N, 'This matrix is not square-like, so it can\'t use PRT to get inverse and you can use MP-inverse to calculate its general inv'
        assert self.Get_row_rank(
            target) == M, 'This matrix is not full-rank, so it can\'t use PRT to get inverse and you can use MP-inverse to calculate its general inv'
        E = np.eye(M)
        A = deepcopy(target)
        if test:
            print('original matrix:')
            print(target)

        for i in range(M):
            for j in range(i, N):
                if test:
                    print('During process in row:{}, col:{}'.format(i, j))
                if sum(abs(A[i:, j])) > eps:
                    # sort rows by non-zero values from small to large
                    zero_index = []
                    non_zero_index = []
                    for k in range(i, M):
                        if abs(A[k, j]) > eps:
                            non_zero_index.append(k)
                        else:
                            zero_index.append(k)

                    non_zero_index = sorted(
                        non_zero_index, key=lambda x: abs(A[x, j]))
                    sort_cols = non_zero_index+zero_index
                    if test:
                        print("Sort_cols index: {}".format(sort_cols))

                    # E should be processed with the same operations.
                    A[i:, :] = A[sort_cols, :]
                    E[i:, :] = E[sort_cols, :]

                    if test:
                        print('After resorting')
                        print(A)
                        print(E)

                    # eliminate elements belows
                    prefix = -A[i+1:, j]/A[i, j]
                    temp = (np.array(prefix).reshape(-1, 1)
                            )@A[i, :].reshape((1, -1))
                    A[i+1:, :] += temp

                    temp = (np.array(prefix).reshape(-1, 1)
                            )@E[i, :].reshape((1, -1))
                    E[i+1:, :] += temp

                    if test:
                        print('After primary row transformation:')
                        print(A)
                        print(E)

                    # eliminate elements above
                    for k in range(i):
                        if abs(A[k, j]) > eps:
                            # be careful to using intermediate parameter
                            temp = -(A[k, j]/A[i, j])
                            A[k, :] += temp*A[i, :]
                            E[k, :] += temp*E[i, :]

                    if test:
                        print('After eliminating elements above:')
                        print(A)
                        print(E)

                    break

                else:
                    continue

        for m in range(M):
            if abs(A[m, m]) > eps:
                # normalize the first element of each row to 1
                if A[m, m] != 1:
                    temp = 1./A[m, m]
                    A[m, :] *= temp
                    E[m, :] *= temp

        return E

    @classmethod
    def Left_inv(self, target):
        assert len(np.shape(target)) == 2
        # get row ans col numbers
        M, N = np.shape(target)
        assert self.Get_col_rank(target) == N
        return self.Primary_Row_Transformation_Inv(target.T@target)@target.T

    @classmethod
    def Right_inv(self, target):
        assert len(np.shape(target)) == 2
        # get row ans col numbers
        M, N = np.shape(target)
        assert self.Get_row_rank(target) == M
        return target.T@self.Primary_Row_Transformation_Inv(target@target.T)

    @classmethod
    def Moore_Penrose_General_Inv(self, target):
        """Moore_Penrose_General_Inv 

        Args:
            target ([np.darray]): [target matrix]

        Returns:
            [np.darray]: [A+]
        """
        assert len(np.shape(target)) == 2
        # get row ans col numbers
        M, N = np.shape(target)
        B, C = self.Full_Rank_Fact(target)
        return self.Right_inv(C)@self.Left_inv(B)

    @classmethod
    def Match_Points_LSM(self, ps, order):
        """Match Points using LSM and show results

        Args:
            ps ([type]): [arrays of (x,y) points]
            order ([type]): [fitting order]
        """
        M, N = ps.shape[0], order+1
        A = np.zeros((M, N))
        # generate coefficient matrix A with given order
        for c in range(N):
            A[:, c] = np.power(ps[:, 0], c)
        # get MP-inv of A
        L_Inv = self.Moore_Penrose_General_Inv(A)
        # solve beta
        beta = L_Inv@ps[:, 1]

        # show results
        print('The matching line: y=', sep='', end='')
        for i in range(order+1):
            print('{:+.3f}*x^{}'.format(beta[i], i), sep='', end='')
        print('\n')
        # calculate fitting error
        error = np.linalg.norm(A@beta-ps[:, 1])
        print("Norm of Matching Error: {:.5f}".format(error))

        # prepare to show matching line
        x = np.arange(np.min(ps[:, 0])-1,
                      np.max(ps[:, 0])+1, 0.01).reshape((-1,))
        B = np.zeros((x.shape[0], N))
        for c in range(N):
            B[:, c] = np.power(x, c)
        y = B@beta

        plt.figure(0)
        fig, ax = plt.subplots()
        raw = ax.scatter(ps[:, 0], ps[:, 1], marker='o', c='blue', s=10)
        match = ax.scatter(x, y, marker='v', c='red', s=1)
        plt.legend([raw, match], ['raw', 'match'], loc='upper left')
        plt.title('Fitting points with {} order'.format(order))
        plt.xlabel('x')
        plt.ylabel('y')
        plt.show()

    @classmethod
    def Gram_Schmidt_Orthogonalization(self, target):
        """Gram_Schmidt_Orthogonalization

        Args:
            target ([np.darray]): [target matrix]

        Returns:
            [np.darray]: [orthogonal matrix]

        Last edited: 22-01-13
        Author: Junno
        """
        # col vectors of target are to be transforms
        assert len(target.shape) == 2, 'shape of target should be (M,N)'
        M, N = target.shape
        # matrix of epsilon
        orth = np.zeros_like(target)
        for i in range(N):
            temp = target[:, i]
            Delta = 0
            for j in range(i):
                Delta -= np.sum(temp*orth[:, j])*orth[:, j]
            temp += Delta
            orth[:, i] = temp/np.linalg.norm(temp)
        return orth

    @classmethod
    def Householder_Transfrom(self, target, project_axis=0, eps=1e-6):
        """Householder_Transfrom

        Args:
            target ([np.darray]): [target matrix]
            project_axis (int, optional): [projection axis]. Defaults to 0.
            eps ([float], optional): [numerical threshold]. Defaults to 1e-6.

        Returns:
            [np.darray]: [the result of householer transform]

        Last edited: 22-01-13
        Author: Junno

        create elementary-reflection-matrix: I-2ww^T, wwT=1
        refer: https://blog.csdn.net/xfijun/article/details/109464005
        """
        # flatten
        x = target.reshape(len(target), 1)
        # generate unit matrix
        unit = np.zeros_like(x, dtype=target.dtype)
        unit[project_axis] = 1
        y = np.linalg.norm(x, ord=2)*unit
        if np.sum(np.abs(x+np.sign(x[project_axis, 0])*y)) < eps:
            return np.eye(len(x))
        else:
            w = (x+np.sign(x[project_axis, 0])*y) / \
                np.linalg.norm(x+np.sign(x[project_axis, 0])*y)
            H = np.eye(len(x))-2*w@w.T
            return H

    @classmethod
    def QR_Fact(self, target, mode='householder'):
        """QR_Fact

        Args:
            target ([np.darray]): [target matrix]
            mode (str, optional): [mode of solver]. Defaults to 'householder'.

        Returns:
            [np.darray]: [Q and R]

        Last edited: 22-01-13
        Author: Junno
        refer: https://blog.csdn.net/xfijun/article/details/109464005
        """
        assert len(target.shape) == 2
        M, N = target.shape
        assert self.Check_full_rank(
            target), "only full_col_rank matrix has QR Fact"
        A = deepcopy(target)
        Q = np.eye(M, dtype=target.dtype)
        if mode == 'householder':
            for i in range(M-1):
                # generate intermediate matrix
                Qi = np.eye(M, dtype=target.dtype)
                # get i-th col
                x = A[i:, i]
                Hi = self.Householder_Transfrom(x)
                # operate householder transform to A
                A[i:, i:] = Hi@A[i:, i:]
                Qi[i:, i:] = Hi
                # A(n-1)=H(n-2)H(n-3)...H(0)A
                Q = Qi@Q
        # induce unit-invertible-matrix D to keep diag elements of R be positive
        D = np.diag(np.where(np.diag(A) < 0, -1, 1))
        R = D@A
        # Q=Q^(-1)D@^(-1)=Q.T@D
        Q = Q.T@D
        return Q, R

    @classmethod
    def Hessenberg(self, target):
        M = target.shape[0]
        Ri = target.copy()
        Q = np.eye(M, dtype=target.dtype)
        for i in range(1, M-1):
            x = Ri[i:, i-1]
            Hi = self.Householder_Transfrom(x)
            Qi = np.eye(M)
            Qi[i:, i:] = Hi
            Q = Qi@Q
            Ri = Qi@Ri@Qi
        return Ri, Q

    @classmethod
    def UR_Fact_Schmidt(self, target):
        """UR_Fact_Schmidt

        Args:
            target ([np.darray]): [target matrix]

        Returns:
            [np.darray]: [U and R]

        Last edited: 22-01-13
        Author: Junno
        """
        assert len(target.shape) == 2
        M, N = target.shape
        assert self.Check_full_rank(
            target), "only invertible matrix has UR Fact"
        U = self.Gram_Schmidt_Orthogonalization(target, norm=True)
        # A=UR,R=U^(-1)A=U^(T)A
        R = U.T@target
        return U, R

    @classmethod
    def det(self, target, eps=1e-6, test=False, np_check=False):
        """calculate determination of matrix: simplify the matrix to the upper triangular matrix

        Args:
            target ([np.darray]): [target M]
            eps ([float]): numerical threshold.
            test (bool, optional): [show checking information]. Defaults to False.
            np_check (bool, optional): [show answer by numpy method to check]. Defaults to False.

        Returns:
            [float]: [det]

        Last edited: 22-01-13
        Author: Junno
        """
        assert len(target.shape) == 2
        M, N = target.shape
        assert M == N

        # fast example
        if M == 1:
            return target[0, 0]
        elif M == 2:
            return target[0, 0]*target[1, 1]-target[0, 1]*target[1, 0]

        # M>=3
        A = deepcopy(target)
        ans = 1
        for i in range(M-1):
            if test:
                print('During process in row:{}, col:{}'.format(i, i))
            if sum(abs(A[i:, i])) > eps:
                first_zero_ind = -1
                first_non_zero_ind = -1
                # resort row
                for k in range(i, M):
                    if abs(A[k, i]) < eps and first_zero_ind < 0:
                        first_zero_ind = k
                        if first_non_zero_ind >= 0:
                            break
                    elif abs(A[k, i]) > eps and first_non_zero_ind < 0:
                        first_non_zero_ind = k
                        if first_zero_ind >= 0:
                            break

                if first_zero_ind == 0:
                    ans *= -1
                    A[[first_zero_ind, first_non_zero_ind], :] = A[[
                        first_non_zero_ind, first_zero_ind], :]

                prefix = -A[i+1:, i]/A[i, i]
                temp = (np.array(prefix).reshape(-1, 1)
                        )@A[i, :].reshape((1, -1))
                A[i+1:, :] += temp

                # calculate det with diag elements
                ans *= A[i, i]

            else:
                return 0

            if test:
                print(A)

        ans *= A[-1, -1]

        if np_check:
            print("numpy result: ", np.linalg.det(target))
            print("my result: ", ans)

        return ans

    @classmethod
    def eig_QR(self, target, mode='householder', epoch=20, eps=1e-4, rnd=3, test=False):
        """eig_QR use QR_Fact to get all eig_values of target M
        *reference: https://blog.csdn.net/xfijun/article/details/109464005
        *通过不断迭代至对角线元素基本保持不变，或者只剩下2*2方块时停止，但是很耗时，暂时没有实现双位移的方法
        *此实现方法对于利用上Hessenberg矩阵化简的优化过程没有明显体现
        *size > 10*10 还是考虑用np的内置方法

        Args:
            target ([float or complex]]): [matrix to get eig_values]
            mode (str, optional): [transform methods]. Defaults to 'householder'.
            epoch (int, optional): [thershold to stop loops]. Defaults to 20.
            eps ([type], optional): [thershold to stop loops]. Defaults to 1e-4.
            test (bool, optional): [show checking information]. Defaults to False.
        """

        assert len(target.shape) == 2
        M, N = target.shape
        assert M == N
        cnt = 0
        # one loop for QR_Fact_Eig

        def QR_epoch(A):
            nonlocal cnt, epoch, mode, eps
            Ak = deepcopy(A)
            flag = 1
            while flag:
                Ak0 = Ak.copy()
                Q, R = self.QR_Fact(Ak, mode)
                Ak = R@Q
                if np.linalg.norm(np.diag(np.abs(Ak-Ak0))) <= eps:
                    flag = 0
                cnt += 1
                if cnt > epoch:
                    flag = 0
            return Ak

        Ak = QR_epoch(target)
        cnt += epoch

        # check result to reloop or stop
        iter_flag = True
        while iter_flag:
            check_complex_eig = False
            eig_list = []

            for i in range(M):
                if check_complex_eig:
                    check_complex_eig = False
                    continue

                # check 2*2 block to get roots or continue QR_epoch loops
                if (i+1) < M and abs(Ak[i+1, i]) >= eps:
                    if i+2 < M and (abs(np.sum(Ak[i+2:, i])) >= eps or abs(np.sum(Ak[i+2:, i+1])) >= eps):
                        Ak = QR_epoch(Ak)
                        break
                    else:
                        check_complex_eig = True
                        if test:
                            print(Ak[i:i+2, i:i+2])
                        c_eig = self.solve_complex_root_QR_eig(
                            Ak[i:i+2, i:i+2])
                        eig_list += c_eig

                else:
                    eig_list.append(Ak[i, i])

            # stop when get to the last row or last two rows is a 2*2 block
            if i == M-1 or (i == M-2 and check_complex_eig == True):
                iter_flag = False

        if test:
            print("Iteration times: ", cnt)
            print(Ak)

        # sort
        eig_list = sorted(
            eig_list, key=lambda x: np.linalg.norm(x), reverse=True)
        return np.round(eig_list, rnd)

    @staticmethod
    def solve_eig_vec(target, eig, dtype=np.complex64, eps=1e-3, test=False):
        """solve_eig_vec: solve according eig_vector given an eig

        Args:
            target ([float or complex]]): [matrix to get eig_vec]
            eig ([float or complex]): [one eig of target]
            dtype ([type], optional): [description]. Defaults to np.complex64.
            eps ([float]): numerical threshold.
            test (bool, optional): [show checking information]. Defaults to False.

        Returns:
            [np.array]: [eig_vector]
        """
        M, N = target.shape
        A = target-np.eye(M, dtype=dtype)*eig
        if test:
            print(A)
        for i in range(M):
            if test:
                print('During process in row:{}, col:{}'.format(i, i))
            if sum(abs(A[i:, i])) > eps:
                zero_index = []
                non_zero_index = []
                for k in range(i, M):
                    if abs(A[k, i]):
                        non_zero_index.append(k)
                    else:
                        zero_index.append(k)

                non_zero_index = sorted(
                    non_zero_index, key=lambda x: abs(A[x, i]))
                sort_cols = non_zero_index+zero_index
                if test:
                    print("Sort_cols index: {}".format(sort_cols))

                A[i:, :] = A[sort_cols, :]

                if test:
                    print('After resorting')
                    print(A)

                if np.sum(abs(A[i+1:, :])):
                    prefix = -A[i+1:, i]/A[i, i]
                    temp = (np.array(prefix).reshape(-1, 1)
                            )@A[i, :].reshape((1, -1))
                    A[i+1:, :] += temp

                # eliminate elements with the same col index above
                # if i != 0:
                #     for j in range(i):
                #         A[j, :] += -(A[j, i]/A[i, i])*A[i, :]

                if test:
                    print('After primary row transformation:')
                    print(A)

            else:
                continue

        # 如果秩不是n-1，即有k重根，需要另外计算，否则所有的特征值给出的特征向量均是一样的，

        # the equation must be non-full-rank, the last row will be left
        for m in range(M-1):
            for n in range(m, N):
                if abs(A[m, n]) > eps:
                    A[m, :] *= 1./A[m, n]
                    break

        if test:
            print(A)

        x = np.zeros((1, M), dtype=dtype)
        for i in range(M-1, -1, -1):
            if i < M-1:
                if np.abs(A[i, i]) < eps:
                    x[0, i] = 0
                else:
                    x[0, i] = -np.sum(A[i, i+1:]*x[0, i+1:])/A[i, i]
            else:
                x[0, i] = 1
            if test:
                print(i, x)

        return x

    @classmethod
    def eig_vv(self, target, eps=1e-3, test=False):
        """eig_vv 求解矩阵的所有特征值与特征向量

        Args:
            target ([float or complex]): [target matrix]
            eps ([float]): numerical threshold.

        Returns:
            [list and array]: [list of eigs ans array of eig_vectors]
        """
        assert target.shape[0] <= 10, 'too large matrix to solve using this imperfect method, it will cost much of time, please use numpy methods instead'
        eigs = self.eig_QR(target, test=test)
        num_eigs = len(eigs)
        eig_vecs = None
        for i in range(num_eigs):
            if test:
                print('solve eigvv')
                print(eigs[i], self.solve_eig_vec(target, eigs[i], eps=eps))
            if i == 0:
                eig_vecs = self.solve_eig_vec(target, eigs[i], eps=eps)
            else:
                eig_vecs = np.concatenate(
                    (eig_vecs, self.solve_eig_vec(target, eigs[i], eps=eps)), axis=0)

        return eigs, eig_vecs

    @classmethod
    def svd(self, target, dtype=np.complex64, eps=1e-4, test=True):
        """svd

        Args:
            target ([float or complex]]): [matrix to get eig_vec]
            dtype ([type], optional): [description]. Defaults to np.complex64.
            eps ([float]): numerical threshold.
            test (bool, optional): [show checking information]. Defaults to False.

        Returns:
            [u,v,d]: [svd components]
        """
        import cmath
        assert len(target.shape) == 2
        M, N = target.shape
        AHA = np.conj(target).T@target
        if test:
            print('AHA:')
            print(AHA)
        eigs, eig_vecs = self.eig_vv(AHA, eps=eps, test=test)
        if test:
            print(eigs, '\n', eig_vecs)
        u_v = None
        temp = False
        for j in range(eig_vecs.shape[0]):
            eig_vecs[j, :] /= np.linalg.norm(eig_vecs[j, :])

            # u_i=(1/sigma_i)*(A@eigv_i) if sigma_i>0  sigma_i=1/sqrt(eig_i)
            if np.linalg.norm(eigs[j]) > eps:
                sigma_j = cmath.sqrt(eigs[j])
                if not temp:
                    u_v = 1/sigma_j*(target@eig_vecs[j, :]).reshape((-1, 1))
                    temp = True
                else:
                    u_v = np.concatenate(
                        (u_v, 1/sigma_j*(target@eig_vecs[j, :]).reshape((-1, 1))), axis=1)
        if test:
            print(eigs, '\n', eig_vecs)
            print(u_v)

        # expand orthogonal basis
        if u_v.shape[1] < M:
            ortho_index = []
            for i in range(u_v.shape[0]):
                if sum(abs(u_v[i, :])) < eps:
                    ortho_index.append(i)
            for i in range(M-u_v.shape[1]):
                u_v = np.concatenate((u_v, np.array([0 if k != ortho_index[i] else 1 for k in range(
                    u_v.shape[0])]).reshape((-1, 1))), axis=1)
            if test:
                print(u_v)

        Sigma = np.zeros((M, N), dtype=dtype)
        for i in range(M):
            if eigs[i] > eps:
                Sigma[i, i] = cmath.sqrt(eigs[i])
            else:
                Sigma[i, i] = eigs[i]

        if test:
            print(Sigma)

        return u_v, Sigma, eig_vecs

    @staticmethod
    def v_L2_norm(vector):
        return np.linalg.norm(vector)

    @staticmethod
    def v_L1_norm(vector):
        return np.sum(np.abs(vector))

    @staticmethod
    def v_infinite_norm(vector):
        return np.max(np.abs(vector))

    @staticmethod
    def v_F_norm(target):
        return np.sqrt(np.trace(target.T@target))

    @staticmethod
    def m_L1_norm(M):
        return np.max(np.sum(np.abs(M), axis=0))

    @staticmethod
    def m_infinite_norm(M):
        return np.max(np.sum(np.abs(M), axis=1))

    @classmethod
    def m_L2_norm(self, M):
        return np.sqrt(np.max(self.eig_QR(np.conj(M).T@M)))

    @classmethod
    def matrix_minimal_polynomial(target):
        return

    @staticmethod
    def solve_square_root(a, b, c):
        temp = b**2-4*a*c
        return [(-b+np.sqrt(temp+0j))/(2*a), (-b-np.sqrt(temp+0j))/(2*a)]

    @classmethod
    def solve_complex_root_QR_eig(self, target):
        return self.solve_square_root(1, -target[0, 0]-target[1, 1], target[0, 0]*target[1, 1]-target[0, 1]*target[1, 0])

    @staticmethod
    def domain_row_transform(arr, eps=1e-6, Test=False):
        """domain_row_transform 选取主元并进行行变换

        Args:
            arr ([np.darray]): [input matrix]
            eps ([float]): numerical threshold.
            test (bool, optional): [show checking information]. Defaults to False.

        Returns:
            [A]: [resort one of arr]
        """
        assert len(arr.shape) == 2
        A = np.copy(arr)
        M, N = A.shape
        for i in range(M):
            for j in range(i, N):
                if Test:
                    print('During process in row:{}, col:{}'.format(i, j))
                if sum(abs(A[i:, j])) > eps:
                    zero_index = []
                    non_zero_index = []
                    for k in range(i, M):
                        if abs(A[k, j]) > eps:
                            non_zero_index.append(k)
                        else:
                            zero_index.append(k)

                    non_zero_index = sorted(
                        non_zero_index, key=lambda x: abs(A[x, j]))
                    sort_idnex = non_zero_index+zero_index

                    if Test:
                        print('Sorting index for {} cols:\n'.format(i), sort_idnex)

                    if sort_idnex[0] != i:
                        A[[sort_idnex[0], i], :] = A[[i, sort_idnex[0]], :]
                    if Test:
                        print('After resorting\n', A)
        return A

    @classmethod
    def Doolittle_Decomposition(self, target, domain_selection=True, test=False):
        """Doolittle_Decomposition equal to LU_Decomposition

        Args:
            target ([np.darray]): [input matrix]
            domain_selection (bool, optional): [operate domain selection on target]. Defaults to True.
            test (bool, optional): [show testing information]. Defaults to False.

        Returns:
            [resort input, L, U] if domain_selection=True else [L,U]

        Author: Junno

        Date: 2022-04-28
        """
        assert len(target.shape) == 2
        arr = np.copy(target)
        M, N = arr.shape
        L = np.zeros_like(arr, dtype=arr.dtype)
        U = np.zeros_like(arr, dtype=arr.dtype)
        if domain_selection:
            arr = self.domain_row_transform(arr)
        for i in range(M):
            if test:
                print(i)
                print(L, U, sep='\n')

            for j in range(i, N):
                if test:
                    print('U', i, j, arr[i, j], L[i, :i], U[:i, j])
                U[i, j] = arr[i, j]-L[i, :i]@U[:i, j]
                if i == j:
                    L[i, j] = 1
                else:
                    if test:
                        print('L', j, i, arr[j, i],
                              L[j, :i], U[:i, i], U[i, i])
                    L[j, i] = (arr[j, i]-(L[j, :i]@U[:i, i]))/U[i, i]
        if domain_selection:
            print('return resort input matrix, L, U')
            return arr, L, U
        else:
            return L, U

    @classmethod
    def Thomas_Linear_Equation_Solve(self, a, b, c, d, eps=1e-4, test=False):
        """Thomas_Linear_Equation_Solve 

        Args:
            a ([list or np.darray]): [lower-secondary diagonal]
            b ([list or np.darray]): [diagonal]
            c ([list or np.darray]): [upper-secondary diagonal]
            d ([list or np.darray]): [b of Ax=b]
            eps ([type], optional): [threshold]. Defaults to 1e-4.
            test (bool, optional): [show testing information]. Defaults to False.

        Returns:
            x [list]: [solution of Ax=b]

        Author: Junno

        Date: 2022-04-28
        """
        assert len(a) == len(c) and len(a)+1 == len(b)
        if not self.check_SDD_vec(a, b, c, eps=eps):
            print(
                'The answer may be wrong because the input are not meet SDD condition completely')
        N = len(b)
        if test:
            print(a, b, c, d)
        beta = np.zeros_like(b)
        x = np.zeros_like(b)
        y = np.zeros_like(b)
        beta[0] = c[0]/b[0]
        y[0] = d[0]/b[0]
        for i in range(1, N):
            if i < N-1:
                beta[i] = c[i]/(b[i]-a[i-1]*beta[i-1])
            y[i] = (d[i]-a[i-1]*y[i-1])/(b[i]-a[i-1]*beta[i-1])
        for i in range(N-1, -1, -1):
            if i == N-1:
                x[i] = y[i]
            else:
                x[i] = y[i]-beta[i]*x[i+1]

        return x

    @staticmethod
    def check_SDD_vec(a, b, c, eps=1e-4, test=False):
        """check_SDD_vec 

        Args:
            a ([np.darray]): [lower-secondary diagonal]
            b ([np.darray]): [diagonal]
            c ([np.darray]): [upper-secondary diagonal]
            eps ([type], optional): [threshold]. Defaults to 1e-4.
            test (bool, optional): [show testing information]. Defaults to False.

        Returns:
            [bool]: [whether or not]

        Author: Junno

        Date: 2022-04-28
        """
        assert len(a) == len(c) and len(a)+1 == len(b)
        N = len(b)
        a_ = np.concatenate(([eps], a), axis=0)
        c_ = np.concatenate((c, [eps]), axis=0)
        if test:
            print(a_, c_)
        sum_a_c = np.abs(a_)+np.abs(c_)
        b_ = np.abs(b)
        for i in range(N):
            if b_[i] < sum_a_c[i]:
                return False
        return True

    @classmethod
    def check_SDDM(self, arr, eps=1e-4):
        """check_SDDM 判断输入矩阵是否满足严格对角占优条件
        Args:
            arr ([np.darray]): [input matrix]
            eps ([type], optional): [threshold]. Defaults to 1e-4.

        Returns:
            [bool]: [whether or not]

        Author: Junno

        Date: 2022-04-28
        """
        assert arr.shape[0] == arr.shape[1]
        return self.check_SDD_vec(np.diagonal(arr, offset=-1), np.diagonal(arr, offset=0), np.diagonal(arr, offset=1), eps=eps)

    @classmethod
    def Serial_iteration(self, A, b, N=0, eps=1e-4, x_0=None, iter_way='Gauss_Seidel', norm='infty', own_method=False, test=False):
        """Serial_iteration:Jacobi and Gauss_Seidel

        Args:
            A ([np.darray]): [A]
            b ([np.darray]): [b]
            N (int, optional): [iteration time, it will give a suitable one when N is 0]. Defaults to 0.
            eps ([float], optional): [error threshold]. Defaults to 1e-4.
            x_0 ([np.darray], optional): [initial x for iteration]. Defaults to None.
            norm (str, optional): [iteration method,'Gauss_Seidel','Jacobi']. Defaults to 'Gauss_Seidel'.
            norm (str, optional): [norm method to calculate fitting error,'infty','L1']. Defaults to 'infty'.
            own_method (bool, optional): [use own method to calculate inv or norm, etc]. Defaults to False.
            test (bool, optional): [show testing information]. Defaults to False.

        Raises:
            ValueError: [The spetrum of Iteration Matrix must be smaller than 1]

        Returns:
            [x]: [solution of Ax=b]

        Author: Junno

        Date: 2022-04-29
        """
        assert A.shape[0] == b.shape[0]
        # use np method or your own
        Inv = np.linalg.inv if not own_method else self.Primary_Row_Transformation_Inv
        SVD = np.linalg.svd if not own_method else self.svd
        # substract D,L,U from A
        D, L, U = np.diag(np.diag(A)), np.tril(A, k=-1), np.triu(A, k=1)
        # calculate B and f
        if iter_way == 'Jacobi':
            B = -Inv(D)@(L+U)
            f = Inv(D)@b
        elif iter_way == 'Gauss_Seidel':
            B = -Inv(D+L)@U
            f = Inv(D+L)@b
        # initial x for iteration
        if x_0 == None:
            x_0 = np.zeros_like(b)

        # get eig of B
        u, s, v = SVD(B)
        # spectrum of B
        eig_B = np.diag(s)
        max_eig_B = np.max(eig_B)
        # check convergence condition
        if max_eig_B > 1:
            print('The spetrum of Iteration Matrix must be smaller than 1')
            raise ValueError

        # matrix norm
        if norm == 'infty':
            F = self.m_infinite_norm
        elif norm == 'L1':
            F = self.m_L1_norm

        q = F(B)
        if test:
            print('q =', q)

        # iteration times
        N = int(np.ceil(np.log10(1/eps)*np.log(10) /
                (-np.log(max_eig_B)))) if N == 0 else N

        # iteration step
        while N:
            x = B@x_0+f
            if q < 1 and q/(1-q)*F(x-x_0) < eps:
                print('last iteration error: ', q/(1-q)*F(x-x_0))
                break
            x_0 = x
            N -= 1

        if test:
            print(
                "The speed of iteration is:{:.2f}".format(-np.log(max_eig_B)))

        return x


if __name__ == "__main__":
    np.set_printoptions(precision=2, suppress=True, linewidth=100)

    # 验证svd分解
    # test_m = np.random.randint(1, 10, (12, 12)).astype(np.float32)
    test_m = np.array([
        [1j, 0, 1j],
        [1, 0, 1],
    ]).reshape((2, 3)).astype(np.complex)
    print('original matrix:\n', test_m)

    # print("numpy solution:")
    # print(np.linalg.svd(test_m))

    # AHA = test_m.T@test_m
    # eigs = Matrix_Solutions.eig_QR(AHA)
    # print('eigs: ', *eigs)

    u, v, d = Matrix_Solutions.svd(test_m, test=True)
    print("SVD Decomposition Result:")
    print('U:\n', u)
    print('Sigma:\n', v)
    print('V:\n', d)

    re = u@v@d
    print('reconstruct target: \n', re)
    print('reconstruct error: ', np.linalg.norm(re-test_m))
    print('reconstruct error with round: ',
          np.linalg.norm(np.round(re, 2)-test_m))

    # a = np.array([[1, 1, 0], [0, 1, 1]]
    #              ).reshape((2, 3)).astype(np.float32)
    # print(a.T)
    # print(Matrix_Solutions.Gram_Schmidt_Orthogonalization(a.T))
    # norm_orth_a = np.linalg.norm()

    # test PRM
    # a = np.array([
    #     [-2, -1, 0, 0],
    #     [0, -1, 0, 0],
    #     [4, 0, 0, -1],
    #     [0, 4, 0, 1]
    # ]).reshape((4, -1)).astype(np.float32)
    # b = np.array([0, 2, 5, -9]).reshape((4, -1)).astype(np.float32)
    # x = Matrix_Solutions.Solve_Ax_b_Equation(a, b, test=True)
    # print(x)
